import tushare as ts
import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt
from datetime import datetime, timedelta

# 初始化pro接口
pro = ts.pro_api('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')

# 1. 查找"天山铝业"股票代码
df_stocks = pro.stock_basic(exchange='', list_status='L')
dongshan = df_stocks[df_stocks['name'].str.contains('天山铝业')]

if dongshan.empty:
    print("未找到名为'天山铝业'的股票")
    exit()

ts_code = dongshan.iloc[0]['ts_code']
symbol = dongshan.iloc[0]['symbol']
name = dongshan.iloc[0]['name']
print(f"找到股票: {name}, 代码: {ts_code}({symbol})")

# 2. 获取3-5年历史行情数据
end_date = datetime.now().strftime('%Y%m%d')
start_date = (datetime.now() - timedelta(days=5 * 365)).strftime('%Y%m%d')

# 获取日线数据
df_daily = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)

# 按日期排序
df_daily = df_daily.sort_values('trade_date')
df_daily['trade_date'] = pd.to_datetime(df_daily['trade_date'])
df_daily.set_index('trade_date', inplace=True)

# 3. 计算技术指标
# 计算移动平均线
df_daily['MA5'] = df_daily['close'].rolling(window=5).mean()
df_daily['MA10'] = df_daily['close'].rolling(window=10).mean()
df_daily['MA20'] = df_daily['close'].rolling(window=20).mean()
df_daily['MA60'] = df_daily['close'].rolling(window=60).mean()

# 计算MACD
exp12 = df_daily['close'].ewm(span=12, adjust=False).mean()
exp26 = df_daily['close'].ewm(span=26, adjust=False).mean()
df_daily['MACD'] = exp12 - exp26
df_daily['MACD_signal'] = df_daily['MACD'].ewm(span=9, adjust=False).mean()
df_daily['MACD_hist'] = df_daily['MACD'] - df_daily['MACD_signal']

# 计算RSI
delta = df_daily['close'].diff()
gain = (delta.where(delta > 0, 0)).rolling(window=14).mean()
loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean()
rs = gain / loss
df_daily['RSI'] = 100 - (100 / (1 + rs))

# 计算布林带
df_daily['STD20'] = df_daily['close'].rolling(window=20).std()
df_daily['UpperBand'] = df_daily['MA20'] + (df_daily['STD20'] * 2)
df_daily['LowerBand'] = df_daily['MA20'] - (df_daily['STD20'] * 2)

# 计算ATR(真实波幅)
df_daily['TR'] = np.maximum(
    df_daily['high'] - df_daily['low'],
    np.maximum(
        abs(df_daily['high'] - df_daily['close'].shift(1)),
        abs(df_daily['low'] - df_daily['close'].shift(1))
    )
)
df_daily['ATR14'] = df_daily['TR'].rolling(window=14).mean()

# 计算成交量均线
df_daily['VMA5'] = df_daily['vol'].rolling(window=5).mean()
df_daily['VMA10'] = df_daily['vol'].rolling(window=10).mean()

# 4. 准备建模数据
# 计算未来N天的收益率 (这里预测未来5天的收益率)
n_days = 5
df_daily['future_return'] = (df_daily['close'].shift(-n_days) / df_daily['close'] - 1) * 100

# 删除包含NaN的行
df_model = df_daily.dropna()

# 特征选择
features = ['open', 'high', 'low', 'close', 'vol',
            'MA5', 'MA10', 'MA20', 'MA60',
            'MACD', 'MACD_signal', 'MACD_hist',
            'RSI', 'UpperBand', 'LowerBand',
            'ATR14', 'VMA5', 'VMA10']

X = df_model[features]
y = df_model['future_return']

# 数据标准化
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# 划分训练集和测试集 (80%训练，20%测试)
X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.2, shuffle=False)

# 5. 构建随机森林模型
rf_model = RandomForestRegressor(
    n_estimators=100,  # 树的数量
    max_depth=10,  # 树的最大深度
    min_samples_split=5,  # 分裂所需最小样本数
    random_state=42
)

# 训练模型
rf_model.fit(X_train, y_train)

# 6. 模型预测
y_pred = rf_model.predict(X_test)

# 7. 模型评估
mse = mean_squared_error(y_test, y_pred)
mae = mean_absolute_error(y_test, y_pred)
r2 = r2_score(y_test, y_pred)

print("\n模型评估结果:")
print(f"均方误差(MSE): {mse:.4f}")
print(f"平均绝对误差(MAE): {mae:.4f}")
print(f"R平方(R2): {r2:.4f}")

# 8. 特征重要性分析
feature_importance = pd.DataFrame({
    'Feature': features,
    'Importance': rf_model.feature_importances_
}).sort_values('Importance', ascending=False)

print("\n特征重要性排序:")
print(feature_importance)

# 9. 可视化预测结果
plt.figure(figsize=(14, 7))
plt.plot(df_model.index[-len(y_test):], y_test, label='实际收益率', color='blue')
plt.plot(df_model.index[-len(y_test):], y_pred, label='预测收益率', color='red', linestyle='--')
plt.title(f'{name} ({ts_code}) 未来{n_days}天收益率预测 vs 实际')
plt.xlabel('日期')
plt.ylabel(f'未来{n_days}天收益率(%)')
plt.legend()
plt.grid(True)
plt.show()

# 10. 模拟投资策略
# 根据预测收益率进行买卖决策
capital = 100000  # 初始资金10万元
position = 0  # 持仓数量
transaction_cost = 0.0015  # 交易成本(买卖各0.15%)
portfolio_value = [capital]

for i in range(len(y_pred)):
    current_price = df_model.iloc[-len(y_test) + i]['close']

    # 当预测收益率大于阈值时买入
    if y_pred[i] > 1.0:  # 预测收益率大于1%
        if position == 0:  # 如果没有持仓，则买入
            position = (capital * (1 - transaction_cost)) / current_price
            capital = 0

    # 当预测收益率小于阈值时卖出
    elif y_pred[i] < -0.5:  # 预测收益率小于-0.5%
        if position > 0:  # 如果有持仓，则卖出
            capital = position * current_price * (1 - transaction_cost)
            position = 0

    # 计算当前投资组合价值
    if position > 0:
        current_value = position * current_price
    else:
        current_value = capital
    portfolio_value.append(current_value)

# 计算策略收益率
strategy_return = (portfolio_value[-1] / portfolio_value[0] - 1) * 100
buy_hold_return = (df_model.iloc[-1]['close'] / df_model.iloc[-len(y_test)]['close'] - 1) * 100

print("\n投资策略表现:")
print(f"策略最终价值: {portfolio_value[-1]:.2f}元")
print(f"策略总收益率: {strategy_return:.2f}%")
print(f"买入持有策略收益率: {buy_hold_return:.2f}%")

# 绘制投资组合价值曲线
plt.figure(figsize=(14, 7))
plt.plot(df_model.index[-len(y_test):], portfolio_value[1:], label='策略价值', color='green')
plt.title(f'{name} ({ts_code}) 投资策略表现')
plt.xlabel('日期')
plt.ylabel('投资组合价值(元)')
plt.legend()
plt.grid(True)
plt.show()